import torch
import torchvision.transforms as T
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
import numpy as np
from PIL import Image
from sklearn.decomposition import PCA

patch_h = 28
patch_w = 28
feat_dim = 768

transform = T.Compose([
    T.GaussianBlur(9, sigma=(0.1, 2.0)),
    T.Resize((patch_h * 14, patch_w * 14)),
    T.CenterCrop((patch_h * 14, patch_w * 14)),
    T.ToTensor(),
    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])

dinov2_vitb14 = torch.hub.load('./model', 'dinov2_vitb14', source='local').cuda()

features_b14 = torch.zeros(4, patch_h * patch_w, feat_dim)
imgs_tensor_b14 = torch.zeros(4, 3, patch_h * 14, patch_w * 14).cuda()

img_path = f'./images/penguin.jpg'
img = Image.open(img_path).convert('RGB')
imgs_tensor_b14[0] = transform(img)[:3]
with torch.no_grad():
    features_dict_b14 = dinov2_vitb14.forward_features(imgs_tensor_b14)
    features_b14 = features_dict_b14['x_norm_patchtokens']

features_b14 = features_b14.reshape(4 * patch_h * patch_w, feat_dim).cpu()
pca = PCA(n_components=3)
pca.fit(features_b14)
pca_features_b14 = pca.transform(features_b14)
pca_features_b14[:, 0] = (pca_features_b14[:, 0] - pca_features_b14[:, 0].min()) / (
        pca_features_b14[:, 0].max() - pca_features_b14[:, 0].min())

pca_features_fg_b14 = pca_features_b14[:, 0] > 0.3
pca_features_bg_b14 = ~pca_features_fg_b14

b = np.where(pca_features_bg_b14)
pca.fit(features_b14[pca_features_fg_b14])
pca_features_rem_b14 = pca.transform(features_b14[pca_features_fg_b14])
for i in range(3):
    pca_features_rem_b14[:, i] = (pca_features_rem_b14[:, i] - pca_features_rem_b14[:, i].min()) \
                                 / (pca_features_rem_b14[:, i].max() - pca_features_rem_b14[:, i].min())

pca_features_rgb_b14 = pca_features_b14.copy()
pca_features_rgb_b14[pca_features_fg_b14] = pca_features_rem_b14
pca_features_rgb_b14[b] = 0
pca_features_rgb_b14 = pca_features_rgb_b14.reshape(4, patch_h, patch_w, 3)
plt.imshow(pca_features_rgb_b14[0][..., ::-1])
plt.savefig('./output/features_b14.png')
plt.show()
plt.close()
print('---s14---')
print(features_b14)
print('---维度---')
print(features_b14.shape)
print('---pca_features---')
print(pca_features_b14)
print('---维度---')
print(pca_features_b14.shape)
